import tkinter as tk
from tkinter import filedialog, ttk
from PIL import Image, ImageTk, ImageEnhance
import cv2
import numpy as np

class ImageEnhancerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("图像画质增强工具")
        self.root.geometry("900x700")
        self.root.configure(bg="#f0f0f0")

        # 创建主框架
        self.main_frame = ttk.Frame(self.root, padding="20")
        self.main_frame.pack(fill=tk.BOTH, expand=True)

        # 文件选择按钮
        self.select_btn = ttk.Button(self.main_frame, text="选择图片", command=self.select_image)
        self.select_btn.pack(pady=10)

        # 功能按钮框架
        self.buttons_frame = ttk.LabelFrame(self.main_frame, text="增强功能", padding="10")
        self.buttons_frame.pack(fill=tk.X, pady=10)

        # 创建按钮容器，使按钮间距更均匀
        self.buttons_container = ttk.Frame(self.buttons_frame)
        self.buttons_container.pack(fill=tk.X, padx=5, pady=5)

        # 画质增强按钮
        self.enhance_btn = ttk.Button(self.buttons_container, text="画质增强", command=self.enhance_image)
        self.enhance_btn.pack(side=tk.LEFT, padx=5)

        # 景物优化按钮
        self.scene_btn = ttk.Button(self.buttons_container, text="景物优化", command=self.scene_enhancement)
        self.scene_btn.pack(side=tk.LEFT, padx=5)

        # 人物增强按钮
        self.face_btn = ttk.Button(self.buttons_container, text="人物增强", command=self.face_enhancement)
        self.face_btn.pack(side=tk.LEFT, padx=5)

        # 锐化按钮
        self.sharpen_btn = ttk.Button(self.buttons_container, text="锐化", command=self.sharpen_image)
        self.sharpen_btn.pack(side=tk.LEFT, padx=5)

        # 对比度调整按钮
        self.contrast_btn = ttk.Button(self.buttons_container, text="对比度调整", command=self.adjust_contrast)
        self.contrast_btn.pack(side=tk.LEFT, padx=5)

        # 亮度调整按钮
        self.brightness_btn = ttk.Button(self.buttons_container, text="亮度调整", command=self.adjust_brightness)
        self.brightness_btn.pack(side=tk.LEFT, padx=5)

        # 重置按钮
        self.reset_btn = ttk.Button(self.buttons_container, text="重置", command=self.reset_image)
        self.reset_btn.pack(side=tk.RIGHT, padx=5)

        # 预览框架
        self.preview_frame = ttk.LabelFrame(self.main_frame, text="图片预览", padding="10")
        self.preview_frame.pack(fill=tk.BOTH, expand=True, pady=10)

        # 预览画布
        self.canvas = tk.Canvas(self.preview_frame, bg="#ffffff")
        self.canvas.pack(fill=tk.BOTH, expand=True)

        # 保存按钮
        self.save_btn = ttk.Button(self.main_frame, text="保存图片", command=self.save_image)
        self.save_btn.pack(pady=10)

        # 初始化变量
        self.original_image = None
        self.processed_image = None
        self.tk_image = None

    def select_image(self):
        file_path = filedialog.askopenfilename(
            filetypes=[("Image files", "*.jpg *.jpeg *.png *.bmp")]
        )
        if file_path:
            # 打开图像
            self.original_image = Image.open(file_path)
            self.processed_image = self.original_image.copy()
            self.display_image()

    def display_image(self):
        # 获取画布尺寸
        canvas_width = self.canvas.winfo_width()
        canvas_height = self.canvas.winfo_height()

        # 如果画布还没有尺寸，使用默认值
        if canvas_width == 1 and canvas_height == 1:
            canvas_width = 800
            canvas_height = 400

        # 调整图像大小以适应画布
        image = self.processed_image.copy()
        img_width, img_height = image.size

        # 计算调整后的尺寸
        ratio = min(canvas_width / img_width, canvas_height / img_height)
        new_width = int(img_width * ratio)
        new_height = int(img_height * ratio)

        # 调整图像大小
        resized_image = image.resize((new_width, new_height), Image.LANCZOS)

        # 转换为Tkinter可用的图像
        self.tk_image = ImageTk.PhotoImage(resized_image)

        # 清除画布并显示新图像
        self.canvas.delete("all")
        x = (canvas_width - new_width) // 2
        y = (canvas_height - new_height) // 2
        self.canvas.create_image(x, y, anchor=tk.NW, image=self.tk_image)

    def enhance_image(self):
        if self.processed_image is None:
            return

        # 转换为OpenCV格式
        cv_image = cv2.cvtColor(np.array(self.processed_image), cv2.COLOR_RGB2BGR)

        # 使用非局部均值去噪
        denoised = cv2.fastNlMeansDenoisingColored(cv_image, None, 10, 10, 7, 21)

        # 锐化
        kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
        sharpened = cv2.filter2D(denoised, -1, kernel)

        # 转换回PIL格式
        self.processed_image = Image.fromarray(cv2.cvtColor(sharpened, cv2.COLOR_BGR2RGB))
        self.display_image()

    def scene_enhancement(self):
        if self.processed_image is None:
            return

        # 转换为OpenCV格式
        cv_image = cv2.cvtColor(np.array(self.processed_image), cv2.COLOR_RGB2BGR)

        # 1. 直方图均衡化（针对彩色图像）
        lab = cv2.cvtColor(cv_image, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=4.0, tileGridSize=(6,6))  # 增加clipLimit，减小tileGridSize
        cl = clahe.apply(l)
        limg = cv2.merge((cl,a,b))
        enhanced = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)

        # 2. 增加饱和度
        hsv = cv2.cvtColor(enhanced, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)
        s = cv2.add(s, 30)  # 增加饱和度
        hsv_enhanced = cv2.merge((h, s, v))
        enhanced = cv2.cvtColor(hsv_enhanced, cv2.COLOR_HSV2BGR)

        # 3. 锐化处理
        kernel = np.array([[0, -1, 0], [-1, 6, -1], [0, -1, 0]])  # 更加强烈的锐化核
        sharpened = cv2.filter2D(enhanced, -1, kernel)

        # 转换回PIL格式
        self.processed_image = Image.fromarray(cv2.cvtColor(sharpened, cv2.COLOR_BGR2RGB))
        self.display_image()

    def face_enhancement(self):
        if self.processed_image is None:
            return

        # 转换为OpenCV格式
        cv_image = cv2.cvtColor(np.array(self.processed_image), cv2.COLOR_RGB2BGR)

        # 加载人脸检测器
        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

        # 转换为灰度图以检测人脸
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.1, 4)  # 调整参数以提高检测率

        # 对检测到的每个人脸进行增强
        for (x, y, w, h) in faces:
            # 提取人脸区域
            face = cv_image[y:y+h, x:x+w]

            # 1. 轻微磨皮效果（高斯模糊）
            blurred_face = cv2.GaussianBlur(face, (3, 3), 0)

            # 2. 锐化处理
            kernel = np.array([[0, -1, 0], [-1, 6, -1], [0, -1, 0]])  # 更强烈的锐化
            sharpened_face = cv2.filter2D(blurred_face, -1, kernel)

            # 3. 调整对比度和亮度
            lab_face = cv2.cvtColor(sharpened_face, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab_face)
            clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8))
            cl = clahe.apply(l)
            limg_face = cv2.merge((cl,a,b))
            enhanced_face = cv2.cvtColor(limg_face, cv2.COLOR_LAB2BGR)

            # 将增强后的人脸放回原图像
            cv_image[y:y+h, x:x+w] = enhanced_face

        # 转换回PIL格式
        self.processed_image = Image.fromarray(cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
        self.display_image()

    def sharpen_image(self):
        if self.processed_image is None:
            return

        enhancer = ImageEnhance.Sharpness(self.processed_image)
        self.processed_image = enhancer.enhance(2.0)
        self.display_image()

    def adjust_contrast(self):
        if self.processed_image is None:
            return

        enhancer = ImageEnhance.Contrast(self.processed_image)
        self.processed_image = enhancer.enhance(1.5)
        self.display_image()

    def adjust_brightness(self):
        if self.processed_image is None:
            return

        enhancer = ImageEnhance.Brightness(self.processed_image)
        self.processed_image = enhancer.enhance(1.2)
        self.display_image()

    def reset_image(self):
        if self.original_image is not None:
            self.processed_image = self.original_image.copy()
            self.display_image()

    def save_image(self):
        if self.processed_image is None:
            return

        file_path = filedialog.asksaveasfilename(
            defaultextension=".jpg",
            filetypes=[("JPEG files", "*.jpg"), ("All files", "*.*")]
        )
        if file_path:
            self.processed_image.save(file_path, "JPEG")

if __name__ == "__main__":
    root = tk.Tk()
    app = ImageEnhancerApp(root)
    root.mainloop()